from mmseg.ops import resize
import torch.nn as nn
import torch.nn.functional as F
import torch

def dual_softmax_loss(X, Y, R,temp = 0.2):
    if X.size() != Y.size() or X.dim() != 2 or Y.dim() != 2:
        raise RuntimeError('Error: X and Y shapes must match and be 2D matrices')

    dist_mat = (X @ Y.t())/0.01
    conf_matrix12 = F.log_softmax(dist_mat, dim=1)
    conf_matrix21 = F.log_softmax(dist_mat.t(), dim=1)
    with torch.no_grad():
        conf12 = torch.exp( conf_matrix12 ).max(dim=-1)[0]
        conf21 = torch.exp( conf_matrix21 ).max(dim=-1)[0]
        conf = conf12 * conf21
    target = torch.arange(len(X), device=X.device)
    score = conf_matrix21+conf_matrix12
    loss = (-score[target, target]*R).mean()


    # loss = F.nll_loss(conf_matrix12, target,weight=R.float()[:,0]) + \
    #        F.nll_loss(conf_matrix21, target)

    # loss = R*loss
    return loss, conf

def compute_accuracy(pred, target, topk=1, thresh=None, ignore_index=None):
    """Calculate accuracy according to the prediction and target.

    Args:
        pred (torch.Tensor): The model prediction, shape (N, num_class, ...)
        target (torch.Tensor): The target of each prediction, shape (N, , ...)
        ignore_index (int | None): The label index to be ignored. Default: None
        topk (int | tuple[int], optional): If the predictions in ``topk``
            matches the target, the predictions will be regarded as
            correct ones. Defaults to 1.
        thresh (float, optional): If not None, predictions with scores under
            this threshold are considered incorrect. Default to None.

    Returns:
        float | tuple[float]: If the input ``topk`` is a single integer,
            the function will return a single float as accuracy. If
            ``topk`` is a tuple containing multiple integers, the
            function will return a tuple containing accuracies of
            each ``topk`` number.
    """
    assert isinstance(topk, (int, tuple))
    if isinstance(topk, int):
        topk = (topk, )
        return_single = True
    else:
        return_single = False

    maxk = max(topk)
    if pred.size(0) == 0:
        accu = [pred.new_tensor(0.) for i in range(len(topk))]
        return accu[0] if return_single else accu
    assert pred.ndim == target.ndim + 1
    assert pred.size(0) == target.size(0)
    assert maxk <= pred.size(1), \
        f'maxk {maxk} exceeds pred dimension {pred.size(1)}'
    pred_value, pred_label = pred.topk(maxk, dim=1)
    # transpose to shape (maxk, N, ...)
    pred_label = pred_label.transpose(0, 1)
    correct = pred_label.eq(target.unsqueeze(0).expand_as(pred_label))
    if thresh is not None:
        # Only prediction values larger than thresh are counted as correct
        correct = correct & (pred_value > thresh).t()
    if ignore_index is not None:
        correct = correct[:, target != ignore_index]
    res = []
    eps = torch.finfo(torch.float32).eps
    for k in topk:
        # Avoid causing ZeroDivisionError when all pixels
        # of an image are ignored
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) + eps
        if ignore_index is not None:
            total_num = target[target != ignore_index].numel() + eps
        else:
            total_num = target.numel() + eps
        res.append(correct_k.mul_(100.0 / total_num))
    return res[0] if return_single else res

def normalize_coords(coord, h, w):
    '''
    turn the coordinates from pixel indices to the range of [-1, 1]
    :param coord: [..., 2]
    :param h: the image height
    :param w: the image width
    :return: the normalized coordinates [..., 2]
    '''
    c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).to(coord.device).float()
    # print(coord[:,:,0].max(), coord[:,:,1].max(), w, h)
    coord_norm = (coord - c) / c
    # print(coord_norm[:,:,0].max(), coord_norm[:,:,1].max(), coord_norm[:,:,0].min(), coord_norm[:,:,1].min())
    return coord_norm

def sample_feat_by_coord(x, coord_n, norm=False):
    '''
    sample from normalized coordinates
    :param x: feature map [batch_size, n_dim, h, w]
    :param coord_n: normalized coordinates, [batch_size, n_pts, 2]
    :param norm: if l2 normalize features
    :return: the extracted features, [batch_size, n_pts, n_dim]
    '''
    feat = F.grid_sample(x, coord_n.unsqueeze(2), padding_mode='zeros', align_corners=False).squeeze(-1)
    # print(feat.shape)
    if norm:
        feat = F.normalize(feat, p=2, dim=1)
    feat = feat.transpose(1, 2)
    return feat


def homogenize(coord):
    # coord = torch.cat((coord, torch.ones_like(coord[:, :, [0]])), -1)
    coord = torch.cat((coord, torch.ones_like(coord[..., [0]])), -1)
    return coord